package org.example.ratelimit;

import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
public class LeakBucketLimiter {

    /**
     * 上次放行时间
     */
    private static long lastTime = System.currentTimeMillis();

    /**
     * 剩余的水量
     */
    private static long water = 0;

    /**
     * 桶的大小
     */
    private static final long capacity = 10;

    /**
     * 每秒放行2个
     */
    private static final long rate = 5;

    /**
     * 以每秒5个的速度进行处理
     *
     * @param taskId
     * @param turn
     * @return
     */
    public static synchronized boolean tryAccquire(long taskId, int turn) {
        long nowTime = System.currentTimeMillis();
        long pastTime = nowTime - lastTime;
        // 以恒定的速率滴水
        water = water - (pastTime * rate) / 1000;
        if (water < 0) {
            water = 0;
        }

        // 请求进入，看剩余的水量是否满了，未满，+1
        if (water < capacity) {
            water++;
            lastTime = nowTime;
            log.info("新时间区到了，taskId：{}, turns：{}", taskId, turn);
            return true;
        } else {
            return false;
        }
    }

    public static void main(String[] args) throws Exception {
        int maxConnections = 100;
        ExecutorService service = Executors.newFixedThreadPool(10);
        CountDownLatch latch = new CountDownLatch(10);

        // 模拟流量激增
        for (int i = 0; i < maxConnections; i++) {
            service.submit(() -> {
                try {
                    long taskId = Thread.currentThread().getId();
                    if (tryAccquire(taskId, 1)) {
                        log.info("线程：{}获取到令牌，放行", taskId);
                        Thread.sleep(200);
                    } else {
                        log.info("线程：{}失败，立马返回",taskId);
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                } finally {
                    latch.countDown();
                }
            });
        }

        latch.await();
    }
}
